#include "trm/reference_frame.h"
#include "trm/span_arithm.h"
#include "trm/mincut.h"
#include "segm/Segmentation.h"
#include "refine/bit_mask.h"
#include "refine/refine2.h"
//#include "refine/refine3.h"
//#include "entropic_tracker/process.h"
#include <imglib/downsample.h>

void Transform(const affine & a, const img::bitmap & src, img::bitmap & tgt); //defined in mask_transform.cpp

namespace trm {

reference_frame::reference_frame(segment & segm) :
	target_frame(segm),
	_segmentation(0)
{
}

reference_frame::~reference_frame()
{
	delete _segmentation;
}

void reference_frame::create(direction dir, span const & computed_span)
{
	//positive time relative segment's origin
	int target = segm().segtime(out_of_span(computed_span, dir)); 
	assert (target > 0);

	unsigned period = segm().param().reference_frame_pitch;
	unsigned ref_time = (target - 1) / period * period;
	if (time() == ref_time)
		return;
	target_frame::create(dir, ref_time);
	prepare();

	if (_segmentation) {
		unsigned last_computed_frame = last_of_span(computed_span, dir);
		twomasks const & last_data = *segm().frame_masks(last_computed_frame);
		if (last_data.ref_frame == ref_time)
			_segmentation->displaceByTransform(last_data.motion);
	}
}

void reference_frame::become(target_frame & t)
{
	swap(t);
	prepare();
}

void reference_frame::prepare()
{
	param const & p = segm().param();

	if (p.downsample > 1) {
		mswin::timing_mark mark(timing_mark("mask downsampling"));
		img::downsample(mask(), _downsampled_mask, p.downsample);
	} else
		_downsampled_mask.clear();

	if (_segmentation) {
		delete _segmentation;
		_segmentation = 0;
	}
	
	if (p.optic_flow) {
		mswin::timing_mark mark(timing_mark("block segmenting"));
		_segmentation = new Segmentation();
		_segmentation->squareCreate(downsampled_mask(), p.tracker_size);
		_segmentation->setSourceImage(&float_image());
		if (_segmentation->count() < 10) {
			delete _segmentation;
			_segmentation = 0;
		}
	}
}

void reference_frame::track_mask(twomasks & m, target_frame const & t)
{
	param const & p = segm().param();

	if (_segmentation) {
		mswin::timing_mark mark(timing_mark("profiles building"));

		ProfileParams pp;
		pp.distinction = DIST_CORRELATION;
		pp.limitativeFunc = LIMF_SMOOTH;
		pp.profileType = PROFILE_AUTOCORRELATION;
		pp.maxDisplacement = p.max_motion;
		pp.maxSelfDisplacement = 2;
		pp.foreignAreaWidth = 10;
		pp.factorSqSigma = 2;
		pp.foreignRho = -1;
		
		_segmentation->buildProfiles(t.float_image(), pp);
	}

	if (_segmentation) {
		mswin::timing_mark mark(timing_mark("coherent motion"));
		affine & rough = m.motion;
		m.ref_frame = time();

		MotionParams mp;
		mp.tau = 1;
		mp.maxIterations = 30;
		mp.minRelativeFimpr = 0.01f;
		mp.muSource = MU_EACH_ITER;
		mp.muValue = 1000;
		mp.muProfileRatio = 0.8f;
		mp.transformClass = TRANSFORM_AFFINE;
		mp.correctionLoops = 2;
		mp.finalCorrection = true;

		_segmentation->coherentMotion(rough, mp);

		affine upsampled_motion(rough);
		upsampled_motion.x *= p.downsample;
		upsampled_motion.y *= p.downsample;
		m.b.resize(mask().dim());
		::Transform(upsampled_motion, mask(), m.b);
	}
	else {
		m.b = mask();
		m.ref_frame = -1;
		m.motion.identity();
	}

	if (p.mincut) {
		mswin::timing_mark mark(timing_mark("mincut segmenting"));
		mincut::params mp;
		mp.pyramid_levels = p.mincut_pyramid_levels;
		mp.max_motion = p.mincut_max_motion;
		mp.hw = p.mincut_max_link_distance;
		mp.nn = p.mincut_link_count;
		mp.aspect = p.mincut_aspect;
		mp.interframe_hw = p.mincut_interframe_max_link_distance;
		mp.interframe_nn = p.mincut_interframe_link_count;
		mp.source = p.mincut_information_source;
		mp.gray_levels_exp = p.mincut_antialiasing_levels_exp;
		mp.timing_tree = &timing_tree();
		trm::grey_mask g;
		mincut::run(image(), mask(),
			t.image(), &m.b, 0, mp, 
			p.mincut_and_antialiasing() ? &g : 0);
		if (!g.empty())
			m.g.store(g);
	}

/*	if (p.entropic_tracker) {
		mswin::timing_mark mark(timing_mark("entropic tracking"));
		colortrack::EntropicTrackerParameters ep;
		ep.s_colorMapper.p_bestMoveIterNum = p.entropic_best_move_iter_num;
		ep.s_colorMapper.p_moveRadius = p.entropic_move_radius;
		ep.s_colorMapper.p_bandRadius = p.entropic_band_radius;
		ep.s_colorMapper.p_unambigThr = p.entropic_unambig_thr;
		ep.s_activeContour.p_alpha = p.entropic_alpha;
		ep.s_activeContour.p_sqrt_control = p.entropic_sqrt_control;
		ep.s_activeContour.p_safe_time_step = p.entropic_safe_time_step;
		EntropicTrack(image(), mask(), t.image(), m.b, ep);
	}*/

	if (p.refine_mask) {
		mswin::timing_mark mark(timing_mark("refining"));
		refine::bit_mask::params rp;
		rp.whs = p.refine_window_half_size;
		rp.auto_t = p.refine_auto;
		rp.t = p.refine_enough_statistics;
		rp.prefer = p.refine_preference;

		affine upsampling_motion(
			p.downsample * m.motion.x,
			p.downsample * m.motion.y,
			p.downsample * m.motion.xx,
			p.downsample * m.motion.xy,
			p.downsample * m.motion.yx,
			p.downsample * m.motion.yy);
		refine::bit_mask::track(
			downsampled_image(), downsampled_mask(),
			t.image(), m.b,
			p.refine_information_source,
			&upsampling_motion, rp);
		if (p.downsample > 1 && p.refine_information_source == refine::source_reference) {
			//m.b seems at this moment blocky, we need second pass to fix it
			affine upsampled_motion(m.motion);
			upsampled_motion.x *= p.downsample;
			upsampled_motion.y *= p.downsample;
			refine::bit_mask::track(
				image(), mask(),
				t.image(), m.b,
				p.refine_information_source,
				&upsampling_motion, rp);
		}
	}

	if (p.refine2_mask) {
		mswin::timing_mark mark(timing_mark("refining2"));
		using namespace refine::refine2;
		params rp;
		rp.max_motion = p.refine2_max_motion;
		rp.max_level = p.refine2_max_level;
		track(image(), mask(), t.image(), m.b, rp);
	}

	/*if (p.refine3_mask) {
		mswin::timing_mark mark(timing_mark("refining3"));
		using namespace refine::refine3;
		params rp;
		rp.t = p.refine3_enough_statistics;
		rp.prefer = p.refine3_preference;

		affine upsampling_motion(
			p.downsample * m.motion.x,
			p.downsample * m.motion.y,
			p.downsample * m.motion.xx,
			p.downsample * m.motion.xy,
			p.downsample * m.motion.yx,
			p.downsample * m.motion.yy);
		track(
			downsampled_image(), downsampled_mask(),
			t.image(), m.b,
			p.refine3_information_source,
			&upsampling_motion, rp);
		if (p.downsample > 1 && p.refine3_information_source == refine::source_reference) {
			//m.b seems at this moment blocky, we need second pass to fix it
			affine upsampled_motion(m.motion);
			upsampled_motion.x *= p.downsample;
			upsampled_motion.y *= p.downsample;
			track(
				image(), mask(),
				t.image(), m.b,
				p.refine3_information_source,
				&upsampling_motion, rp);
		}
	}*/
}

} //namespace trm
